#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <vector>
#include <cmath>
#include <memory>
#include <iostream>
#include <stack> 
#include <queue> 

namespace py = pybind11;

struct Node {
    int id;                             
    double singularValue;               
    std::vector<std::shared_ptr<Node>> neighbors; 
    std::vector<std::shared_ptr<Node>> mutexNodes;
    
    Node(int nodeId, double value) : id(nodeId), singularValue(value) {}

    void addNeighbor(std::shared_ptr<Node> neighbor) {
        neighbors.push_back(neighbor);
    }
    
    void addMutexNode(std::shared_ptr<Node> mutexNode) {
        mutexNodes.push_back(mutexNode);
    }

    void printNode() const {
        std::cout << "Node ID: " << id << ", Singular Value: " << singularValue << std::endl;
        std::cout << "Neighbors: ";
        for(size_t i = 0; i < neighbors.size(); i++){
            std::cout << neighbors[i]->id << " ";
        }
        std::cout << std::endl;
        std::cout << "Mutex Nodes: ";
        for (const auto& mutexNode : mutexNodes) {
            std::cout << mutexNode->id << " ";
        }
        std::cout << std::endl;
    }
};

std::vector<std::shared_ptr<Node>> buildGraph(const std::vector<double>& sValueVec, const std::vector<std::vector<bool>>& connectivityMatrix) {
    std::vector<std::shared_ptr<Node>> nodes;
    for (size_t nodeId = 0; nodeId<sValueVec.size();nodeId++) {
        nodes.push_back(std::make_shared<Node>(nodeId, sValueVec[nodeId]));
    }
    size_t rows = connectivityMatrix.size();
    size_t cols;
    for (size_t i = 0; i < rows; i++) {
        cols = connectivityMatrix[i].size();
        for (size_t j = 0; j < cols; j++) {
            if (connectivityMatrix[i][j]) {
                nodes[i]->addNeighbor(nodes[j]);
            }
            else{
                nodes[i]->addMutexNode(nodes[j]);
            }
        }
    }
    return nodes;
}

struct IndexedValue {
    double value;
    size_t index;
};

bool compareIndexedValue(const IndexedValue& a, const IndexedValue& b) {
    return a.value < b.value;
};

std::vector<std::vector<size_t>> arg_sort_by_row(const std::vector<std::vector<double>>& matrix) {
    std::vector<std::vector<size_t>> sortedIndices;
    for (const auto& row : matrix) {
        std::vector<IndexedValue> indexedRow;
        for (size_t i = 0; i < row.size(); ++i) {
            indexedRow.push_back({row[i], i});
        }
        std::sort(indexedRow.begin(), indexedRow.end(), compareIndexedValue);
        std::vector<size_t> sortedIndicesRow;
        for (const auto& indexedValue : indexedRow) {
            sortedIndicesRow.push_back(indexedValue.index);
        }
        sortedIndices.push_back(sortedIndicesRow);
    }
    return sortedIndices;
}

std::vector<double> singularity(const std::vector<std::vector<double>>& matrix, const std::vector<std::vector<size_t>>& sortedIndices, size_t m) {
    std::vector<double> sums;
    for (size_t i = 0; i < matrix.size(); ++i) {
        double sum = 0.0;
        for (size_t j = 0; j < m && j < sortedIndices[i].size(); j++) {
            sum += matrix[i][sortedIndices[i][j]];
        }
        sums.push_back(sum);
    }
    return sums;
}

std::vector<std::vector<bool>> buildConnectivityMatrix(const std::vector<std::vector<double>>& matrix, const std::vector<std::vector<size_t>>& sortedIndices, size_t m) {
    size_t rows = matrix.size();
    size_t cols = matrix[0].size();
    std::vector<std::vector<bool>> connectivityMatrix(rows, std::vector<bool>(cols, false));
    for (size_t i = 0; i < rows; ++i) {
        std::vector<bool> isMinElement(cols, false);
        for (size_t j = 0; j < m && j < sortedIndices[i].size(); ++j) {
            isMinElement[sortedIndices[i][j]] = true;
        }
        for (size_t j = 0; j < cols; ++j) {
            if (!isMinElement[j]) {
                connectivityMatrix[i][j] = true;
            }
        }
    }
    return connectivityMatrix;
}


std::pair<std::vector<std::vector<bool>>, std::vector<double>> initProblem(const std::vector<std::vector<double>>& matrix, size_t M){
    std::vector<std::vector<size_t>> sortedIndices = arg_sort_by_row(matrix);
    std::vector<double> s_values = singularity(matrix, sortedIndices, M);
    std::vector<std::vector<bool>> connectivityMatrix = buildConnectivityMatrix(matrix, sortedIndices, M);
    return {connectivityMatrix, s_values};
}

/************************ Depth First Search***********************************/
bool isMutex(const std::shared_ptr<Node>& node, const std::unordered_set<int>& visited) {
    for (const auto& mutexNode : node->mutexNodes) {
        if (visited.find(mutexNode->id) != visited.end()) {
            return true;
        }
    }
    return false;
}

std::pair<std::unordered_set<int>, double> DFS(const std::vector<std::shared_ptr<Node>>& graph, int maxDepth) {
    std::stack<std::pair<std::shared_ptr<Node>, int>> stack;
    std::unordered_set<int> visited; // to be optimized;
    std::unordered_set<int> best_visit;
    std::vector<std::shared_ptr<Node>> path;
    for (const auto& startNode : graph) {
        stack.push({startNode, 0});
    }
    
    double pathSum = 0.0;
    double maxSingularitySum = 0.0;
    while (!stack.empty()) {
        auto [node, depth] = stack.top();
        stack.pop();

        while((depth+1) < path.size()){ // step back
            std::shared_ptr<Node> tmp_node = path.back();
            path.pop_back();
            visited.erase(tmp_node->id);
            pathSum -= tmp_node->singularValue;
        }

        if (isMutex(node, visited)) {
            continue;
        }

        if (visited.find(node->id) != visited.end()) {
            continue;
        }

        visited.insert(node->id);
        path.push_back(node);
        pathSum += node->singularValue;

        if (depth + 1 >= maxDepth) { 
            if(pathSum > maxSingularitySum){
                best_visit = visited;
                maxSingularitySum = pathSum;
            }
            
            std::cout << "Current Path: ";
            for (const auto& p : path) {
                std::cout << p->id << " ";
            }
            std::cout << " | Singular Value Sum: " << pathSum << std::endl;

            // step back
            pathSum -= node->singularValue;
            path.pop_back();
            visited.erase(node->id);
            continue;
        }

        // expand the 'stack' of the nodes to visit
        for (const auto& neighbor : node->neighbors) {
            if (visited.find(neighbor->id) == visited.end()) {
                stack.push({neighbor, depth + 1});
            }
        }

    }
    return {best_visit, maxSingularitySum};
}

/******************Beam K Search****************************/

struct Path {
    std::vector<std::shared_ptr<Node>> nodes;
    std::unordered_set<int> visited;
    double singularValueSum;

    Path() : singularValueSum(0.0) {}

    Path(const std::vector<std::shared_ptr<Node>>& n, double sum) : nodes(n), singularValueSum(sum) {}

    bool operator<(const Path& other) const {
        return singularValueSum > other.singularValueSum; // descending order
    }
};

std::pair<std::unordered_set<int>, double> beamSearch(const std::vector<std::shared_ptr<Node>>& graph, int maxDepth, int beamWidth) {
    /*
        beamWidth = 1: Greedy Algorithm;
        beamWidth -> \infty; Breadth-first search
    */
    std::list<Path> beam;
    std::unordered_set<int> visited;
    std::unordered_set<int> best_visit;
    double maxSingularitySum = 0.0;

    for (const auto& startNode : graph) {
        Path initialPath;
        initialPath.nodes.push_back(startNode);
        initialPath.visited.insert(startNode->id);
        initialPath.singularValueSum = startNode->singularValue;
        beam.push_back(initialPath);
    }

    for (int depth = 1; depth < maxDepth; ++depth) {
        beam.sort();
        if (beam.size() > beamWidth) {
            auto it = std::next(beam.begin(), beamWidth);
            beam.erase(it, beam.end());
        }

        std::cout << "Depth " << depth << " Paths:\n";
        for (const auto& path : beam) {
            std::cout << "Path: ";
            for (const auto& node : path.nodes) {
                std::cout << node->id << " ";
            }
            std::cout << "| Singular Value Sum: " << path.singularValueSum << std::endl;
        }

        std::list<Path> nextBeam;
        for (const auto& path : beam) {
            std::shared_ptr<Node> currentNode = path.nodes.back();
            for (const auto& neighbor : currentNode->neighbors) {
                if (path.visited.find(neighbor->id) == path.visited.end() && !isMutex(neighbor, path.visited)) {
                    Path newPath = path;
                    newPath.nodes.push_back(neighbor);
                    newPath.visited.insert(neighbor->id);
                    newPath.singularValueSum += neighbor->singularValue;
                    nextBeam.push_back(newPath);
                }
            }
        }
        beam = nextBeam;
    }

    std::cout << "Depth " << maxDepth << " Paths:\n";
    for (const auto& path : beam) {
        std::cout << "Path: ";
        for (const auto& node : path.nodes) {
            std::cout << node->id << " ";
        }
        std::cout << "| Singular Value Sum: " << path.singularValueSum << std::endl;
    }

    for(const auto& path : beam){
        if(path.singularValueSum > maxSingularitySum){
            maxSingularitySum = path.singularValueSum;
            best_visit = path.visited;
        }
    }
    return {best_visit, maxSingularitySum};
}

std::vector<int> test(const std::vector<std::vector<double>>& matrix, size_t M, size_t K, size_t beamWidth=100) {
    std::pair<std::vector<std::vector<bool>>, std::vector<double>> result = initProblem(matrix, M);
    const auto& connectivityMatrix = result.first;
    const auto& singularValues = result.second;
    std::vector<std::shared_ptr<Node>> graph = buildGraph(singularValues, connectivityMatrix);
    // DFS(graph, K);
    beamWidth = beamWidth < matrix.size() ? beamWidth : matrix.size();
    std::pair<std::unordered_set<int>, double> beam_rst = beamSearch(graph, K, beamWidth);
    const std::unordered_set<int>& best_visit = beam_rst.first;
    return std::vector<int>(best_visit.begin(), best_visit.end());
}

PYBIND11_MODULE(example, m) {
    m.doc() = "A module to discover the K-group casess that have the maximum singularity.";
    m.def("arg_sort_by_row", &arg_sort_by_row, "Sort each row of the matrix, and return indices");
    m.def("singularity", &singularity, "Calculate the sum of the first m elements of each row in the matrix",
          py::arg("matrix"), py::arg("sortedIndices"), py::arg("m"));
    m.def("buildConnectivityMatrix", &buildConnectivityMatrix, "build the Connectivity Matrix",
          py::arg("matrix"), py::arg("sortedIndices"), py::arg("m"));
    m.def("initProblem", &initProblem, "Init the problem with M as the size of neighbour",
          py::arg("matrix"), py::arg("m"));
    m.def("test", &test, "test interface for debugging the library!",
          py::arg("matrix"), py::arg("m"), py::arg("k"), py::arg("beamWidth"));
}
